# demo_fastapi.py
import asyncio
from pathlib import Path
from typing import Optional, Any, Dict, Union
import uvicorn
import numpy as np
import pandas as pd
import nest_asyncio

import pandas as pd
from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, FileResponse
from pydantic import BaseModel

import graphrag.api as api
from graphrag.config.load_config import load_config

# 应用nest_asyncio来支持嵌套事件循环
nest_asyncio.apply()

# 声明FastAPI
app = FastAPI(title="GraphRAG API", description="基于知识图谱的检索增强生成系统")

# 挂载静态文件 - 使用当前目录
app.mount("/static", StaticFiles(directory="."), name="static")


# 请求模型
class QueryRequest(BaseModel):
    query: str
    method: str = "local"  # local, global, drift, basic
    response_type: str = "text"  # text, json
    community_level: int = 1
    dynamic_community_selection: bool = False


def clean_for_json(obj: Any) -> Any:
    """
    清理对象以便JSON序列化，处理numpy数据类型等

    Args:
        obj: 需要清理的对象

    Returns:
        清理后的对象
    """
    if obj is None:
        return None
    elif isinstance(obj, (str, int, float, bool)):
        return obj
    elif isinstance(obj, np.integer):
        return int(obj)
    elif isinstance(obj, np.floating):
        return float(obj)
    elif isinstance(obj, np.bool_):
        return bool(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif isinstance(obj, pd.Series):
        return obj.tolist()
    elif isinstance(obj, pd.DataFrame):
        return obj.to_dict('records')
    elif isinstance(obj, dict):
        return {key: clean_for_json(value) for key, value in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [clean_for_json(item) for item in obj]
    else:
        # 对于其他对象，尝试转换为字符串
        try:
            return str(obj)
        except:
            return None


async def run_graphrag_api(api_func, *args, **kwargs):
    """
    在正确的事件循环中运行GraphRAG API函数
    
    Args:
        api_func: GraphRAG API函数
        *args: 位置参数
        **kwargs: 关键字参数
        
    Returns:
        API函数的返回值
    """
    try:
        # 确保在正确的事件循环中运行
        loop = asyncio.get_event_loop()
        if loop.is_running():
            # 如果当前事件循环正在运行，直接调用
            return await api_func(*args, **kwargs)
        else:
            # 如果事件循环没有运行，创建新的任务
            task = asyncio.create_task(api_func(*args, **kwargs))
            return await task
    except RuntimeError as e:
        if "event loop" in str(e).lower():
            # 事件循环错误，尝试在新的循环中运行
            try:
                return await asyncio.wait_for(api_func(*args, **kwargs), timeout=60)
            except Exception as inner_e:
                raise Exception(f"GraphRAG API执行失败: {inner_e}")
        else:
            raise e


# 全局数据加载（在应用启动时加载一次）
PROJECT_DIRECTORY = "./datas/demo2"
graphrag_config = None
entities = None
communities = None
community_reports = None
text_unit = None
relations_ships = None


@app.on_event("startup")
async def load_data():
    """在应用启动时加载数据"""
    global graphrag_config, entities, communities, community_reports, text_unit, relations_ships

    graphrag_config = load_config(Path(PROJECT_DIRECTORY))

    entities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/entities.parquet")
    communities = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/communities.parquet")
    community_reports = pd.read_parquet(
        f"{PROJECT_DIRECTORY}/output/community_reports.parquet"
    )
    text_unit = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/text_units.parquet")
    relations_ships = pd.read_parquet(f"{PROJECT_DIRECTORY}/output/relationships.parquet")


@app.get("/", response_class=HTMLResponse)
async def read_index():
    """提供主页HTML"""
    # 尝试多个可能的路径
    possible_paths = [
        "index.html",
        "src/index.html",
        "./index.html",
        "./src/index.html"
    ]
    
    for path in possible_paths:
        if Path(path).exists():
            return FileResponse(path)
    
    # 如果文件不存在，返回一个简单的HTML页面
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>GraphRAG API</title>
        <meta charset="UTF-8">
    </head>
    <body>
        <h1>GraphRAG API 服务器运行中</h1>
        <p>API文档: <a href="/docs">/docs</a></p>
        <p>请确保index.html文件存在于正确的位置</p>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)


@app.post("/query")
async def query(request: QueryRequest):
    """
    统一的查询接口，支持所有查询类型

    Args:
        request: 查询请求对象

    Returns:
        dict: 包含响应和上下文的字典
    """
    try:
        # 根据查询类型调用相应的API
        if request.method == "global":
            response, context = await run_graphrag_api(
                api.global_search,
                config=graphrag_config,
                entities=entities,
                communities=communities,
                community_reports=community_reports,
                community_level=request.community_level,
                dynamic_community_selection=request.dynamic_community_selection,
                response_type=request.response_type,
                query=request.query,
            )
        elif request.method == "local":
            response, context = await run_graphrag_api(
                api.local_search,
                config=graphrag_config,
                entities=entities,
                communities=communities,
                community_reports=community_reports,
                community_level=request.community_level,
                text_units=text_unit,
                covariates=None,
                relationships=relations_ships,
                response_type=request.response_type,
                query=request.query,
            )
        elif request.method == "drift":
            # 漂移查询暂时使用本地搜索实现
            response, context = await run_graphrag_api(
                api.local_search,
                config=graphrag_config,
                entities=entities,
                communities=communities,
                community_reports=community_reports,
                community_level=request.community_level,
                text_units=text_unit,
                covariates=None,
                relationships=relations_ships,
                response_type=request.response_type,
                query=request.query,
            )
            # 添加漂移查询的标识
            response = f"[漂移查询模式] {response}"
        elif request.method == "basic":
            # 基础查询暂时使用本地搜索实现
            response, context = await run_graphrag_api(
                api.local_search,
                config=graphrag_config,
                entities=entities,
                communities=communities,
                community_reports=community_reports,
                community_level=request.community_level,
                text_units=text_unit,
                covariates=None,
                relationships=relations_ships,
                response_type=request.response_type,
                query=request.query,
            )
            # 添加基础查询的标识
            response = f"[基础查询模式] {response}"
        else:
            return {"error": f"不支持的查询类型: {request.method}"}

        # 清理数据以便JSON序列化
        cleaned_context = clean_for_json(context)

        # 根据响应类型返回不同格式
        if request.response_type == "json":
            return {
                "response": response,
                "context": cleaned_context,
                "query": request.query,
                "method": request.method
            }
        else:
            return {
                "response": response,
                "context": cleaned_context
            }

    except Exception as e:
        error_msg = str(e)
        if "event loop" in error_msg.lower():
            return {"error": f"事件循环错误: {error_msg}. 请尝试重启服务器。"}
        elif "semaphore" in error_msg.lower():
            return {"error": f"并发控制错误: {error_msg}. 请尝试重启服务器。"}
        else:
            return {"error": f"查询执行失败: {error_msg}"}


@app.get("/global_search")
async def global_search(
    query: str,
    community_level: int = 1,
    response_type: str = "Multiple Paragraphs"
):
    """
    全局搜索接口

    Args:
        query: 搜索查询语句
        community_level: 社区层级
        response_type: 响应类型

    Returns:
        dict: 包含响应和上下文的字典
    """
    try:
        response, context = await run_graphrag_api(
            api.global_search,
            config=graphrag_config,
            entities=entities,
            communities=communities,
            community_reports=community_reports,
            community_level=community_level,
            dynamic_community_selection=False,
            response_type=response_type,
            query=query,
        )
        return {
            "response": response,
            "context": clean_for_json(context)
        }
    except Exception as e:
        return {"error": str(e)}


@app.get("/local_search")
async def local_search(
    query: str,
    community_level: int = 2,
    response_type: str = "Multiple Paragraphs"
):
    """
    本地搜索接口

    Args:
        query: 搜索查询语句
        community_level: 社区层级
        response_type: 响应类型

    Returns:
        dict: 包含响应和上下文的字典
    """
    try:
        response, context = await run_graphrag_api(
            api.local_search,
            config=graphrag_config,
            entities=entities,
            communities=communities,
            community_reports=community_reports,
            community_level=community_level,
            text_units=text_unit,
            covariates=None,
            relationships=relations_ships,
            response_type=response_type,
            query=query,
        )
        return {
            "response": response,
            "context": clean_for_json(context)
        }
    except Exception as e:
        return {"error": str(e)}


@app.get("/drift_search")
async def drift_search(
    query: str,
    community_level: int = 2,
    response_type: str = "Multiple Paragraphs"
):
    """
    漂移搜索接口（暂时使用本地搜索实现）

    Args:
        query: 搜索查询语句
        community_level: 社区层级
        response_type: 响应类型

    Returns:
        dict: 包含响应和上下文的字典
    """
    try:
        response, context = await run_graphrag_api(
            api.local_search,
            config=graphrag_config,
            entities=entities,
            communities=communities,
            community_reports=community_reports,
            community_level=community_level,
            text_units=text_unit,
            covariates=None,
            relationships=relations_ships,
            response_type=response_type,
            query=query,
        )
        return {
            "response": f"[漂移查询模式] {response}",
            "context": clean_for_json(context)
        }
    except Exception as e:
        return {"error": str(e)}


@app.get("/basic_search")
async def basic_search(
    query: str,
    community_level: int = 2,
    response_type: str = "Multiple Paragraphs"
):
    """
    基础搜索接口（暂时使用本地搜索实现）

    Args:
        query: 搜索查询语句
        community_level: 社区层级
        response_type: 响应类型

    Returns:
        dict: 包含响应和上下文的字典
    """
    try:
        response, context = await run_graphrag_api(
            api.local_search,
            config=graphrag_config,
            entities=entities,
            communities=communities,
            community_reports=community_reports,
            community_level=community_level,
            text_units=text_unit,
            covariates=None,
            relationships=relations_ships,
            response_type=response_type,
            query=query,
        )
        return {
            "response": f"[基础查询模式] {response}",
            "context": clean_for_json(context)
        }
    except Exception as e:
        return {"error": str(e)}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

